[v0.9.1] MTP supports V1 scheduler #2371
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the attention metadata handling by introducing a centralized AscendCommonAttentionMetadata dataclass. This is a good architectural improvement that centralizes logic and reduces code duplication. However, the review identified several critical issues related to this refactoring, including incorrect tensor slicing and initialization that could lead to runtime errors or incorrect behavior. Specifically, there are bugs in build_dummy_metadata in attention_v1.py and in the build method of mla_v1.py's metadata builder. There are also some redundant code assignments that should be cleaned up.
| block_table[:num_reqs, :self.runner.max_num_blocks_per_req] = ( | ||
| block_table[:num_reqs]) | ||
|
|
||
| query_start_loc = common_attn_metadata.query_start_loc |
There was a problem hiding this comment.
The common_attn_metadata used here is initialized without a query_start_loc, causing common_attn_metadata.query_start_loc to be None. This None value is then passed to the AscendMetadata constructor, which will lead to a runtime error as query_start_loc is a required tensor. To fix this, a dummy query_start_loc tensor should be created for the decode-only state.
| query_start_loc = common_attn_metadata.query_start_loc | |
| query_start_loc = torch.arange(0, num_reqs + 1, dtype=torch.int32, device=block_table.device) |
vllm_ascend/attention/mla_v1.py
Outdated
| actual_seq_lengths_q = query_start_loc[1:num_decodes+1].tolist() | ||
| max_seq_lens = seq_lens[:num_decodes].max().item() | ||
| seq_lens = seq_lens[:num_decodes] | ||
| input_positions = input_positions[:num_decodes] |
There was a problem hiding this comment.
The input_positions tensor is incorrectly sliced using num_decodes (the number of decode requests). It should be sliced with num_decode_tokens (the number of tokens in decode requests) to select the correct positions for the decode phase. This is a critical bug that will lead to incorrect behavior.
| input_positions = input_positions[:num_decodes] | |
| input_positions = input_positions[:num_decode_tokens] |
| self.device, non_blocking=True) | ||
| attn_mask = common_attn_metadata.attn_mask | ||
| attn_state = common_attn_metadata.attn_state | ||
| query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[:num_reqs + 1] |
There was a problem hiding this comment.
This line is redundant as query_start_loc_cpu is already assigned with the same value on line 183. Please remove this duplicate assignment to improve code clarity and avoid potential confusion.
| input_positions = self.runner.positions_cpu[:num_tokens].to( | ||
| device, non_blocking=True).long() |
There was a problem hiding this comment.
The input_positions tensor is assigned here from a CPU tensor and then immediately overwritten on line 398 with a device tensor from common_attn_metadata.positions. This first assignment is redundant and can be removed for clarity and to avoid unnecessary operations.
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
| ): | ||
| self.vllm_config = vllm_config | ||
| self.model_config = vllm_config.model_config | ||
| self.device = device | ||
| self.runner = runner |
JC-ut0
force-pushed
the
v0.9.1-dev
branch
2 times, most recently
from
619fff4 to
96a8bfb
Compare
github-actions
bot
added
the
documentation
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
github-actions
bot
removed
the
documentation
github-actions
bot
added
the
documentation
JC-ut0
changed the title
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Signed-off-by: xuyexiong <[email protected]>
github-actions
bot
removed
merge-conflicts
documentation
What this PR does / why we need it?
Does this PR introduce any user-facing change?
How was this patch tested?